package com.atguigu.aiproject.vectorstore;

import com.atguigu.aiproject.mapper.CreationMapper;
import com.atguigu.aiproject.vo.VecStoreCreationDetailVO;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.stereotype.Component;

import java.io.File;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

// todo 理论上不应属于业务层
// 向量数据库定制增强，如果由deepseek生成则代码行数与逻辑复杂度指数增长，时间复杂度与空间复杂度暴增
@Component
@Slf4j
public class CreationVectorStoreManager {
    private final CreationMapper creationMapper;
    private final VectorStore creationVectorStore;
    // todo: 向量数据库存储目录
    private final String VECTOR_STORE_DIRECTORY = "vectorstore";
    private final String FILE_SUFFIX = ".vecstojson";
    private final String MAIN_CREATION_NAME = "main" + FILE_SUFFIX;

    public CreationVectorStoreManager(CreationMapper creationMapper, VectorStore creationVectorStore) {
        this.creationMapper = creationMapper;
        this.creationVectorStore = creationVectorStore;
    }

    // 重新从数据库加载所有作品
    public void reloadCreationVectorStore() {
        this.clear();
        this.loadCreationVectorStoreFromDB();
    }

    // 加载向量数据库
    public void loadCreationVectorStore() {
        final File directory = new File(VECTOR_STORE_DIRECTORY);
        final File mainFile = new File(directory, MAIN_CREATION_NAME);
        if(!mainFile.exists() || !mainFile.isFile()){
            this.loadCreationVectorStoreFromDB();
            return;
        }

        final SimpleVectorStore simpleVectorStore = (SimpleVectorStore) creationVectorStore;

        for(File file : directory.listFiles()){
            if(file.getName().endsWith(FILE_SUFFIX)){
                simpleVectorStore.load(file);
            }
        }
    }

    // 完全重置向量数据库
    public void loadCreationVectorStoreFromDB() {
        List<VecStoreCreationDetailVO> creationDetails = creationMapper.selectAllVecStoreCreationDetails();
        // 转换为Document并向量化
        List<Document> documents = new ArrayList<>();

        for (VecStoreCreationDetailVO creationDetail: creationDetails) {
            // 创建包含丰富信息的文档内容
            Document doc = documentTransformer(creationDetail);
            documents.add(doc);
        }

        // 批量添加文档到向量数据库
        creationVectorStore.add(documents);

        // 将向量数据库内容存到本地
        // 将多个文件合并到一个文件中
        final File directory = new File(VECTOR_STORE_DIRECTORY);
        if(!directory.exists()){
            directory.mkdirs();
        }
        if(!directory.isFile()){
            directory.delete();
            directory.mkdirs();
        }
        for(File file : directory.listFiles()){
            if(file.getName().endsWith(FILE_SUFFIX)){
                file.delete();
            }
        }
        final File mainFile = new File(directory, MAIN_CREATION_NAME);
        ((SimpleVectorStore) creationVectorStore).save(mainFile);

        System.out.println("成功向量化 " + creationDetails.size() + " 篇作品");
    }

    // 清空向量数据库
    public void clear(){
        // todo topK应设置为数据库所有作品的数量
        SearchRequest searchRequest = new SearchRequest.Builder().topK(150).build();
        // VectorStore接口过少，只能先获取所有文档再删除，不如c++一根
        List<Document> documents = creationVectorStore.similaritySearch(searchRequest);
        List<String> documentIds = new ArrayList<>();
        for (Document document : documents){
            documentIds.add(document.getId());
        }
        creationVectorStore.delete(documentIds);
    }

    // 删除指定id的作品
    public void delete(int creationId){
        SearchRequest searchRequest = new SearchRequest.Builder().topK(1)
                .filterExpression(String.format("\"作品ID\" == \"%d\"", creationId))
                .build();
        // 正常情况下只有一个结果
        List<Document> documents = creationVectorStore.similaritySearch(searchRequest);
        if(documents.isEmpty()){
            log.warn("删除失败：向量数据库中未找到id为 {} 的作品", creationId);
        }else creationVectorStore.delete(documents.get(0).getId());
    }

    // 从数据库添加一个作品
    public void add(int creationId){
        VecStoreCreationDetailVO creationDetail = creationMapper.selectVecStoreCreationDetailById(creationId);
        Document doc = documentTransformer(creationDetail);
        creationVectorStore.add(List.of(doc));
    }

    // 从数据库更新一个作品
    public void update(int creationId){
        this.delete(creationId);
        this.add(creationId);
    }

    // 添加一个作品
    public void add(VecStoreCreationDetailVO creationDetail){
        Document doc = documentTransformer(creationDetail);
        creationVectorStore.add(List.of(doc));
    }

    // 更新一个作品
    public void update(VecStoreCreationDetailVO creationDetail){
        this.delete(creationDetail.getCreationId());
        this.add(creationDetail);
    }

    // 提取作品信息以向量化
    private Document documentTransformer(VecStoreCreationDetailVO creationDetail) {
        return new Document("内容" + creationDetail.getContent(),
                Map.of(
                        "作品ID", creationDetail.getCreationId().toString(),
                        "标题", creationDetail.getTitle(),
                        "关键词", creationDetail.getKeyword(),
                        "作品分类名", creationDetail.getCgyName(),
                        "作者", creationDetail.getUserNickName()
                )
        );
    }

    public List<Document> similaritySearch(SearchRequest searchRequest) {
        return creationVectorStore.similaritySearch(searchRequest);
    }
}
